# layers/kan_ffn.py
import torch
import torch.nn as nn
from layers.efficient_kan.kan import KANLinear

class KAN_FFN(nn.Module):
    def __init__(self, d_model, d_ff=None, dropout=0.1):
        super().__init__()
        d_ff = d_ff or 4 * d_model
        self.kan = nn.Sequential(
            KANLinear(d_model, d_ff),
            nn.ReLU(),
            KANLinear(d_ff, d_model)
        )
        self.dropout = nn.Dropout(dropout)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x):
        residual = x
        B, L, D = x.shape
        x = self.kan(x.view(B * L, D)).view(B, L, D)
        x = self.dropout(x)
        return self.norm(residual + x)
